from typing import List, Dict, Any, Optional, Tuple
from datetime import datetime, timedelta
from utils.logger import Logger
from utils.exceptions import ToolError, AgentNotFoundError
from utils.constants import TASK_STATUSES, AGENT_STATUSES
from states.agent_state import AgentStateManager
from states.task_state import TaskStateManager
from middleware.redis_adapter import get_redis_adapter
from distributed.distributed_lock import DistributedLock  # 复用分布式锁
from auxiliary_agents.compliance import check_compliance  # 复用合规检查

# 初始化日志
logger = Logger.get_logger(agent_name="orchestrator_tool")


class OrchestratorTools:
    """编排核心工具集：任务拆解、Agent选择、调度、结果整合"""

    def __init__(self):
        """初始化：状态管理器、Redis客户端、分布式锁"""
        try:
            self.agent_state_manager = AgentStateManager()
            self.task_state_manager = TaskStateManager()
            self.redis_client = get_redis_adapter()  # 使用Redis适配器，支持集群模式
            self.distributed_lock = DistributedLock()
            self.max_retries = 3  # 子任务最大重试次数
            self.retry_interval = 2  # 重试间隔（秒）
            self.task_timeout = 30  # 子任务超时时间（秒）
            logger.info("OrchestratorTools initialized successfully")
        except Exception as e:
            logger.error(f"OrchestratorTools初始化失败: {str(e)}", exc_info=True)
            raise ToolError(
                message="编排工具初始化失败",
                context={"error": str(e)}
            ) from e

    def parse_complex_task(self, task: Dict[str, Any]) -> List[Dict[str, Any]]:
        """解析复杂任务为子任务列表（基于规则/LLM）"""
        task_id = task.get("task_id", "unknown")
        task_type = task.get("task_type", "general")
        query = task.get("query", "")

        try:
            # 1. 合规性检查（全局过滤）
            compliance_result = check_compliance(query)
            if not compliance_result["allowed"]:
                raise ToolError(
                    message=f"任务不合规: {compliance_result['reason']}",
                    context={"task_id": task_id, "query": query}
                )

            # 2. 基于规则的任务拆解（示例：支持多模态、工具调用、文档问答的组合任务）
            sub_tasks = []
            if "图片" in query or "截图" in query:
                # 多模态子任务
                sub_tasks.append({
                    "sub_task_id": f"{task_id}_mm_001",
                    "task_type": "multi_modal",
                    "query": query,
                    "metadata": {"priority": "high", "parent_task_id": task_id}
                })
            elif "查询" in query or "天气" in query or "计算" in query:
                # 工具调用子任务（自动识别工具类型）
                tool_name = "weather_query" if "天气" in query else "math_calculator" if "计算" in query else "web_search"
                sub_tasks.append({
                    "sub_task_id": f"{task_id}_tool_001",
                    "task_type": "tool_call",
                    "tool_name": tool_name,
                    "params": self._extract_tool_params(query, tool_name),
                    "metadata": {"priority": "medium", "parent_task_id": task_id}
                })
            elif "文档" in query or "文件" in query:
                # 文档问答子任务
                doc_id = task.get("metadata", {}).get("doc_id")
                if not doc_id:
                    raise ToolError(message="文档问答任务缺少doc_id参数", context={"task_id": task_id})
                sub_tasks.append({
                    "sub_task_id": f"{task_id}_doc_001",
                    "task_type": "doc_qa",
                    "query": query,
                    "metadata": {"priority": "medium", "parent_task_id": task_id, "doc_id": doc_id}
                })
            else:
                # 文本问答子任务
                sub_tasks.append({
                    "sub_task_id": f"{task_id}_text_001",
                    "task_type": "text_qa",
                    "query": query,
                    "metadata": {"priority": "normal", "parent_task_id": task_id}
                })

            logger.debug(f"复杂任务拆解完成: task_id={task_id}, 子任务数={len(sub_tasks)}")
            return sub_tasks
        except ToolError:
            raise
        except Exception as e:
            logger.error(f"任务拆解失败: task_id={task_id}, query={query}, error={str(e)}", exc_info=True)
            raise ToolError(
                message="任务拆解失败",
                context={"task_id": task_id, "query": query, "error": str(e)}
            ) from e

    def _extract_tool_params(self, query: str, tool_name: str) -> Dict[str, Any]:
        """从查询中提取工具参数（基于规则）"""
        params = {}
        if tool_name == "weather_query":
            # 提取城市名（示例："查询北京的天气" → 城市=北京）
            import re
            city_pattern = r"([北京上海广州深圳天津重庆杭州武汉成都南京西安郑州长沙青岛等一二线城市名]|[\u4e00-\u9fa5]{2,4})"
            match = re.search(city_pattern, query)
            if match:
                params["city"] = match.group(1)
        elif tool_name == "math_calculator":
            # 提取数学表达式（示例："计算1+2*3" → expression=1+2*3）
            import re
            expr_pattern = r"[0-9+\-*/()\[\]{}.,sin cos tan log sqrt pow abs round]+"
            matches = re.findall(expr_pattern, query)
            if matches:
                params["expression"] = "".join(matches)
        elif tool_name == "web_search":
            # 提取搜索关键词（整个查询作为关键词）
            params["query"] = query.replace("查询", "").replace("搜索", "").strip()
        return params

    def select_best_agent(self, task_type: str) -> Tuple[str, str]:
        """选择最优Agent实例（基于健康状态、任务数最少）"""
        # Agent类型与服务名映射
        agent_type_map = {
            "text_qa": "text_qa_agent",
            "doc_qa": "doc_qa_agent",
            "multi_modal": "multi_modal_agent",
            "tool_call": "tool_call_agent"
        }

        agent_service_name = agent_type_map.get(task_type)
        if not agent_service_name:
            raise AgentNotFoundError(message=f"无对应Agent类型: {task_type}")

        # 获取所有健康的Agent实例
        healthy_agents = self.agent_state_manager.get_healthy_agents(agent_service_name)
        if not healthy_agents:
            raise AgentNotFoundError(message=f"无健康的{agent_service_name}实例")

        # 选择当前任务数最少的Agent
        best_agent = min(healthy_agents, key=lambda x: x.get("task_count", 0))
        agent_url = best_agent.get("service_url")  # Agent的API地址（如：http://127.0.0.1:8001）
        agent_id = best_agent.get("agent_id")

        if not agent_url:
            raise AgentNotFoundError(message=f"Agent实例缺少服务地址: {agent_id}")

        logger.debug(f"选择最优Agent: task_type={task_type}, agent_id={agent_id}, agent_url={agent_url}")
        return agent_id, agent_url

    def dispatch_sub_task(self, sub_task: Dict[str, Any]) -> Dict[str, Any]:
        """调度子任务到目标Agent（带重试机制）"""
        sub_task_id = sub_task.get("sub_task_id", "unknown")
        task_type = sub_task.get("task_type")
        retries = 0

        while retries < self.max_retries:
            try:
                # 1. 更新子任务状态为"运行中"
                self.task_state_manager.update_task_state(
                    task_id=sub_task_id,
                    status=TASK_STATUSES.RUNNING,
                    message=f"开始执行子任务（重试次数：{retries}）"
                )

                # 2. 选择最优Agent
                agent_id, agent_url = self.select_best_agent(task_type)

                # 3. 发送任务到Agent（调用/predict接口）
                headers = {"Content-Type": "application/json"}
                from apps.tool_call.api_client import APIClient
                api_client = APIClient()
                response = api_client.send_request(
                    url=f"{agent_url}/predict",
                    method="POST",
                    headers=headers,
                    json=sub_task,
                    timeout=self.task_timeout
                )

                # 4. 更新子任务状态为"成功"
                self.task_state_manager.update_task_state(
                    task_id=sub_task_id,
                    status=TASK_STATUSES.COMPLETED,
                    message="子任务执行成功",
                    result=response
                )

                logger.debug(f"子任务调度成功: sub_task_id={sub_task_id}, agent_id={agent_id}")
                return {
                    "sub_task_id": sub_task_id,
                    "agent_id": agent_id,
                    "status": TASK_STATUSES.COMPLETED,
                    "result": response
                }
            except AgentNotFoundError as e:
                # Agent未找到，直接失败（无需重试）
                logger.error(f"子任务调度失败（无可用Agent）: sub_task_id={sub_task_id}, error={str(e)}")
                self.task_state_manager.update_task_state(
                    task_id=sub_task_id,
                    status=TASK_STATUSES.FAILED,
                    message=f"子任务执行失败: {str(e)}"
                )
                return {
                    "sub_task_id": sub_task_id,
                    "status": TASK_STATUSES.FAILED,
                    "error_msg": str(e)
                }
            except Exception as e:
                retries += 1
                logger.warning(f"子任务调度失败（将重试）: sub_task_id={sub_task_id}, retries={retries}, error={str(e)}")
                self.task_state_manager.update_task_state(
                    task_id=sub_task_id,
                    status=TASK_STATUSES.RETRYING,
                    message=f"子任务执行失败，准备重试（重试次数：{retries}）"
                )

                # 重试间隔
                import time
                time.sleep(self.retry_interval)

        # 重试次数耗尽
        error_msg = f"子任务执行失败（重试{self.max_retries}次耗尽）"
        logger.error(f"{error_msg}: sub_task_id={sub_task_id}")
        self.task_state_manager.update_task_state(
            task_id=sub_task_id,
            status=TASK_STATUSES.FAILED,
            message=error_msg
        )
        return {
            "sub_task_id": sub_task_id,
            "status": TASK_STATUSES.FAILED,
            "error_msg": error_msg
        }

    def integrate_sub_task_results(
            self,
            parent_task_id: str,
            sub_task_results: List[Dict[str, Any]],
            original_query: str
    ) -> str:
        """整合所有子任务结果，生成最终答案（调用LLM）"""
        try:
            # 过滤成功的子任务结果
            successful_results = [res for res in sub_task_results if res["status"] == TASK_STATUSES.COMPLETED]
            if not successful_results:
                return f"所有子任务执行失败，无法回答原问题：{original_query}"

            # 构建整合提示词
            results_context = "\n".join([
                f"子任务{idx + 1}（{res['sub_task_id']}）结果：{str(res['result'])[:200]}"
                for idx, res in enumerate(successful_results)
            ])

            prompt = f"""
            原问题：{original_query}
            子任务执行结果：
            {results_context}

            请基于所有子任务的成功结果，整合生成一个简洁、准确、完整的回答，无需提及子任务细节，直接回应原问题。
            如果子任务结果冲突，以最新的结果为准；如果结果不完整，基于已有信息合理补充。
            """

            # 调用LLM生成最终答案
            from apps.tool_call.api_client import APIClient
            api_client = APIClient()
            api_params = {
                "model": "gpt-3.5-turbo",
                "messages": [{"role": "user", "content": prompt}],
                "max_tokens": 1024
            }
            response = api_client.call_openai_api(api_params)
            final_answer = response["choices"][0]["message"]["content"]

            logger.debug(f"子任务结果整合完成: parent_task_id={parent_task_id}")
            return final_answer
        except Exception as e:
            logger.error(f"结果整合失败: parent_task_id={parent_task_id}, error={str(e)}", exc_info=True)
            raise ToolError(
                message="结果整合失败",
                context={"parent_task_id": parent_task_id, "error": str(e)}
            ) from e


# 单例工具实例
orchestrator_tools = OrchestratorTools()

# 测试入口
if __name__ == "__main__":
    test_logger = Logger.update_context(task_id="orchestrator_test", agent_name="orchestrator_tool")
    test_logger.info("开始测试编排工具")

    # 测试复杂任务
    test_task = {
        "task_id": "complex_task_001",
        "task_type": "complex",
        "query": "查询北京的天气，然后计算3的平方，最后总结结果",
        "metadata": {"priority": "high", "user_id": "test_user"}
    }

    try:
        # 1. 任务拆解
        sub_tasks = orchestrator_tools.parse_complex_task(test_task)
        test_logger.info(f"任务拆解结果: {sub_tasks}")

        # 2. 调度子任务
        sub_task_results = []
        for sub_task in sub_tasks:
            result = orchestrator_tools.dispatch_sub_task(sub_task)
            sub_task_results.append(result)
        test_logger.info(f"子任务调度结果: {sub_task_results}")

        # 3. 结果整合
        final_answer = orchestrator_tools.integrate_sub_task_results(
            parent_task_id=test_task["task_id"],
            sub_task_results=sub_task_results,
            original_query=test_task["query"]
        )
        test_logger.info(f"最终整合答案: {final_answer}")

        test_logger.success("编排工具测试完成")
    except ToolError as e:
        test_logger.error(f"测试失败: {e}")